#include <openssl/bn.h>
#include <openssl/ossl_typ.h>
#include <openssl/rand.h>
#include <stdbool.h>
#include <stdio.h>

typedef struct {
	int bits;
	BIGNUM* n;  // n
	BIGNUM* hs; // g
	BIGNUM* n2; // n^2
} paillier_pubkey;

typedef struct {
	int bits;
	BIGNUM* lambda;  // p
	BIGNUM* mu;  // q
} paillier_prikey;

typedef struct {
	paillier_pubkey* pubkey;
	paillier_prikey* prikey;
	int bits;
} paillier_keypair;

// 前向声明
struct bignum_st {
	BN_ULONG* d; /* Pointer to an array of 'BN_BITS2' bit
				  * chunks. */
	int top;     /* Index of last used d +1. */
	/* The next are internal book keeping for bn_expand. */
	int dmax; /* Size of the d array. */
	int neg;  /* one if the number is negative */
	int flags;
};

// lcm算法
BIGNUM* BN_lcm(BIGNUM* r, const BIGNUM* a, const BIGNUM* b, BN_CTX* ctx) {
	bool is_ctx = true;
	if (ctx == NULL) {
		is_ctx = false;
		ctx = BN_CTX_new();
	}
	BIGNUM* gcd = BN_new();
	BIGNUM* tmp_a = BN_dup(a), * tmp_b = BN_dup(b);
	BN_gcd(gcd, a, b, ctx);
	BN_div(tmp_a, NULL, a, gcd, ctx);
	BN_mul(r, tmp_a, tmp_b, ctx);
	BN_free(gcd);
	BN_free(tmp_a);
	BN_free(tmp_b);
	if (!is_ctx) {
		BN_CTX_free(ctx);
	}
	return r;
}

// 生成介于[0, n) 范围内的随机整数
BIGNUM* generate_random_bn(BIGNUM* n) {
	int num_bits = BN_num_bits(n);
	int num_bytes = (num_bits + 7) / 8;

	BIGNUM* rnd = BN_new();
	unsigned char* buf = (unsigned char*)malloc(num_bytes);
	if (!rnd || !buf) {
		// handle error
		return NULL;
	}

	do {
		if (!RAND_bytes(buf, num_bytes)) {
			// handle error
			return NULL;
		}
		BN_bin2bn(buf, num_bytes, rnd);
	} while (BN_cmp(rnd, n) >= 0);

	free(buf);
	return rnd;
}

// 生成介于 0 和 2^k（k 为正整数）之间的随机整数
BIGNUM* generate_random2_bn(int k) {
	BN_CTX* ctx = BN_CTX_new();
	int num_bytes = (k + 7) / 8;
	unsigned char* buf = malloc(num_bytes);
	if (!RAND_bytes(buf, num_bytes)) {
		// Handle error ...
	}

	BIGNUM* bn = BN_bin2bn(buf, num_bytes, NULL);
	if (!bn) {
		// Handle error ...
	}

	// Limit the range of the BN to [0, 2^k)
	BIGNUM* k_bn = BN_new();
	BN_set_bit(k_bn, k);
	BN_mod(bn, bn, k_bn, ctx);

	free(buf);
	BN_free(k_bn);
	BN_CTX_free(ctx);
	return bn;
}

// new paillier_keypair
paillier_keypair* paillier_keypair_new() {
	paillier_keypair* keypair =
		(paillier_keypair*)malloc(sizeof(paillier_keypair));
	if (keypair) {
		keypair->bits = 0;
		keypair->prikey = (paillier_prikey*)malloc(sizeof(paillier_prikey));
		if (keypair->prikey) {
			keypair->prikey->bits = 0;
			keypair->prikey->lambda = BN_new();
			keypair->prikey->mu = BN_new();
		}
		keypair->pubkey = (paillier_pubkey*)malloc(sizeof(paillier_pubkey));
		if (keypair->pubkey) {
			keypair->pubkey->bits = 0;
			keypair->pubkey->hs = BN_new();
			keypair->pubkey->n = BN_new();
			keypair->pubkey->n2 = BN_new();
		}
	}
	return keypair;
}

// free paillier_keypair
void paillier_keypair_free(paillier_keypair* keypair) {
	free(keypair->prikey);
	free(keypair->pubkey);
	free(keypair);
	return;
}

// // 生成p,q对
// void generate_p_q(paillier_keypair *pk, int bits) {
//   BIGNUM *p = pk->p;
//   BIGNUM *q = pk->q;
//   pk->bits = bits;
//   BN_CTX *ctx = BN_CTX_new();
//   // BIGNUM* temp = BN_new();
//   // BIGNUM* four = BN_new();
//   // BIGNUM* three = BN_new();

//   BN_generate_prime_ex(p, bits / 2, 1, NULL, NULL, NULL);
//   BN_ULONG temp = BN_div_word(p, 4);
//   for (; temp != 3;) { // p%4=3
//     BN_generate_prime_ex(p, bits / 2, 1, NULL, NULL, NULL);
//   }
//   BN_copy(q, p);
//   for (; BN_cmp(q, p) == 0;) {
//     BN_generate_prime_ex(q, bits / 2, 1, NULL, NULL, NULL);
//     temp = BN_div_word(q, 4);
//     for (; temp != 3;) { // q%4=3
//       BN_generate_prime_ex(q, bits / 2, 1, NULL, NULL, NULL);
//     }
//   }
//   // BN_free(three);
//   // BN_free(four);
//   // BN_free(temp);
//   BN_CTX_free(ctx);
//   return;
// }

// 生成密钥对
void generate_keypair(paillier_keypair* pk, int bits) {
	BN_CTX* ctx = BN_CTX_new();
	pk->bits = bits;
	pk->pubkey->bits = bits;
	pk->prikey->bits = bits;
	BIGNUM* times = BN_new();
	BIGNUM* p = BN_new();
	BIGNUM* q = BN_new();
	BIGNUM* ps1 = BN_new();
	BIGNUM* qs1 = BN_new();
	BIGNUM* temp = BN_new();
	BIGNUM* three = BN_new();
	BIGNUM* x = BN_new();
	BIGNUM* h = BN_new();

	// 生成p, q
	BN_set_word(times, 4);
	BN_generate_prime_ex(p, bits / 2, 1, NULL, NULL, NULL);
	BN_div(NULL, temp, p, times, ctx); // p%4
	BN_set_word(three, 3);
	for (; BN_cmp(temp, three) != 0;) { // p%4=3
		BN_generate_prime_ex(p, bits / 2, 1, NULL, NULL, NULL);
		BN_div(NULL, temp, p, times, ctx); // p%4
	}
	BN_copy(q, p);
	for (; BN_cmp(q, p) == 0;) {
		BN_generate_prime_ex(q, bits / 2, 1, NULL, NULL, NULL);
		BN_div(NULL, temp, q, times, ctx); // q%4
		for (; BN_cmp(temp, three) != 0;) {            // q%4=3
			BN_generate_prime_ex(q, bits / 2, 1, NULL, NULL, NULL);
			BN_div(NULL, temp, q, times, ctx); // q%4
		}
	}

	BN_set_word(times, 2);

	BN_mul(pk->pubkey->n, p, q, ctx); // n=p*q
	BN_set_word(times, 1);

	BN_sub(ps1, p, times);     // ps1 = (p - 1)
	BN_mul(temp, ps1, q, ctx); // (p - 1) * q

	BN_sub(qs1, q, times);     // qs1 = (q - 1)
	BN_mul(temp, qs1, p, ctx); // (q - 1) * p

	BN_mul(temp, ps1, qs1, ctx); // lambda = (p - 1)(q - 1)
	BN_rshift1(pk->prikey->lambda, temp);    // lambda = lambda除以2
	BN_mod_inverse(pk->prikey->mu, pk->prikey->lambda, pk->pubkey->n,
		ctx); // mu = lambda^-1 mod n

	BN_rand_range(x, pk->pubkey->n); // 生成(0, n)的整数
	BN_set_word(temp, 2);
	BN_mod_exp(h, x, temp, pk->pubkey->n, ctx);       // h = x^2 mod n
	h->neg = 1;                                       // h = -x^2 mod n
	BN_exp(pk->pubkey->n2, pk->pubkey->n, temp, ctx); // n^2
	BN_mod_exp(pk->pubkey->hs, h, pk->pubkey->n, pk->pubkey->n2,
		ctx); // hs = h^n mod n^2

	BN_free(h);
	BN_free(x);
	BN_free(temp);
	BN_free(ps1);
	BN_free(qs1);
	BN_free(q);
	BN_free(p);
	BN_free(times);
	BN_CTX_free(ctx);
	return;
}

int encrypt(BIGNUM* out, paillier_pubkey* pubkey, BIGNUM* m) {
	if (pubkey->bits == 0 || (pubkey->bits & 1) != 0) {
		return 1;
	}
	BN_CTX* ctx = BN_CTX_new();
	BIGNUM* temp = BN_new();
	BIGNUM* temp1 = BN_new();
	BIGNUM* nm = BN_new();
	BIGNUM* one = BN_new();
	BIGNUM* alpha = generate_random2_bn(pubkey->bits / 2);

	BN_mul(nm, pubkey->n, m, ctx);                        // nm = n * m
	BN_add(temp1, nm, one);                               // n * m + 1
	BN_mod_exp(temp, pubkey->hs, alpha, pubkey->n2, ctx); // hs^alpha
	BN_mod_mul(out, temp, temp1, pubkey->n2,
		ctx); // c = (n * m + 1)hs^alphq mod n^2

	BN_free(alpha);
	BN_free(nm);
	BN_free(one);
	BN_free(temp);
	BN_free(temp1);
	BN_CTX_free(ctx);
	return 0;
}

int decrypt(BIGNUM* out, paillier_keypair* pk, BIGNUM* c) {
	BN_CTX* ctx = BN_CTX_new();
	BN_MONT_CTX* mont_ctx = BN_MONT_CTX_new();
	BIGNUM* times = BN_new();
	BIGNUM* temp = BN_new();
	BIGNUM* n2 = BN_new();
	BIGNUM* lnx = BN_new();
	BIGNUM* temp2 = BN_new();

	BN_set_word(times, 2);
	BN_exp(n2, pk->pubkey->n, times, ctx);
	BN_mod_exp_mont(lnx, c, pk->prikey->lambda, pk->pubkey->n, ctx, mont_ctx);  // lnx = c^lambda mod n^2
	//printf("debug: c: %s\t lambda: %s\t n2: %s\t lnx: %s\n", BN_bn2dec(c),
	//	        BN_bn2dec(pk->prikey->lambda), BN_bn2dec(pk->pubkey->n2), BN_bn2dec(lnx) );
	BN_set_word(times, 1);	
	BN_sub(temp, lnx, times);	// lnx - 1
	BN_div(temp2, NULL, temp, pk->pubkey->n, ctx);
	BN_mod_mul(out, temp2, pk->prikey->mu, pk->pubkey->n, ctx);

	BN_free(temp2);
	BN_free(temp);
	BN_free(lnx);
	BN_free(n2);
	BN_free(times);
	BN_MONT_CTX_free(mont_ctx);
	BN_CTX_free(ctx);

	return 0;
}

int main() {
	paillier_keypair* keypair = paillier_keypair_new();
	// generate_p_q(keypair, 1024);
	generate_keypair(keypair, 64);
	printf(" n: %s\n hs: %s\n n2: %s\n lambda: %s\n mu: %s\n",
		BN_bn2dec(keypair->pubkey->n), BN_bn2dec(keypair->pubkey->hs),
		BN_bn2dec(keypair->pubkey->n2), BN_bn2dec(keypair->prikey->lambda),
		BN_bn2dec(keypair->prikey->mu));
	BIGNUM* input = BN_new();
	BIGNUM* output = BN_new();
	BN_set_word(input, 1314);
	printf("明文: %s\n", BN_bn2dec(input));
	encrypt(output, keypair->pubkey, input);
	printf("密文: %s\n", BN_bn2dec(output));
	decrypt(input, keypair, output);
	printf("解密文: %s\n", BN_bn2dec(input));

	BN_free(output);
	BN_free(input);
	paillier_keypair_free(keypair);
}

// void paillier_gen_keypair(paillier_pubkey* pk, int bits) {
//     BN_CTX* ctx = BN_CTX_new();
//     BIGNUM* p = BN_new();
//     BIGNUM* q = BN_new();
//     BIGNUM* p1 = BN_new();
//     BIGNUM* q1 = BN_new();
//     BIGNUM* n = BN_new();
//     BIGNUM* n1 = BN_new();
//     BIGNUM* n2 = BN_new();
//     BIGNUM* g = BN_new();
//     BIGNUM* lambda = BN_new();
//     BIGNUM* mu = BN_new();
//     BIGNUM* temp = BN_new();
//     BN_generate_prime_ex(p, bits / 2, 1, NULL, NULL, NULL);    //
//     随机生成两个质数 BN_generate_prime_ex(q, bits / 2, 1, NULL, NULL, NULL);
//     BN_mul(n, p, q, ctx);   // 计算 n = p*q
//     BN_mul(n2, n, n, ctx);  // 计算 n^2
//     BN_sub(p1 p, BN_value_one());   // p1 = p-1
//     BN_sub(q1, q, BN_value_one());   // q1 = q-1
// }
//
// void paillier_genkey(paillier_keypair* kp, int bits) {
//     BN_CTX* ctx = BN_CTX_new();
//     BIGNUM* p = BN_new();
//     BIGNUM* q = BN_new();
//     BIGNUM* p1 = BN_new();
//     BIGNUM* q1 = BN_new();
//     BIGNUM* n = BN_new();
//     BIGNUM* n1 = BN_new();
//     BIGNUM* n2 = BN_new();
//     BIGNUM* lambda = BN_new();
//     BIGNUM* g = BN_new();
//     BIGNUM* temp = BN_new();
//     BN_generate_prime_ex(p, bits / 2, 1, NULL, NULL, NULL);    //
//     随机生成两个质数 BN_generate_prime_ex(q, bits / 2, 1, NULL, NULL, NULL);
//     BN_mul(n, p, q, ctx);   // 计算 n = p*q
//     BN_mul(n2, n, n, ctx);  // 计算 n^2
//     BN_sub(p1 p, BN_value_one());   // p1 = p-1
//     BN_sub(q1, q, BN_value_one());   // q1 = q-1
//     BN_lcm(lambda, p1, q1, ctx);  // 计算 lambda = lcm(p-1, q-1)
//     BN_add(n1, n, 1);   // 计算n+1
//     BN_set_word(temp, 2);
//     BN_mod_exp(g, n1, temp, n2, ctx);  // 计算 g = (1+n)^2 mod n^2
//     BN_gcd(p1, q, lambda, ctx);  // p = gcd(p-1, q-1)
//     BN_div(lambda, NULL, lambda, p, ctx);  // lambda = lambda/p
//     kp->n = n;
//     kp->g = g;
//     BN_free(p);
//     BN_free(q);
//     BN_free(n2);
//     BN_free(lambda);
//     BN_CTX_free(ctx);
// }
//
// void paillier_encrypt(BIGNUM* c, BIGNUM* m, BIGNUM* r, paillier_keypair* kp)
// {
//     BN_CTX* ctx = BN_CTX_new();
//     BIGNUM* n = kp->n;
//     BIGNUM* n2 = BN_new();
//     BN_mul(n2, n, n, ctx);  // n^2
//     BIGNUM* g = kp->g;
//     BIGNUM* tmp = BN_new();
//     BIGNUM* tmp2 = BN_new();
//     BIGNUM* r2 = BN_new();
//     BIGNUM* c2 = BN_new();
//     RAND_priv_bytes(r, BN_num_bits(n));  // 随机生成一个 r
//     BN_mod_exp(tmp, g, m, n2, ctx);     // g^m mod n^2
//     BN_mod_exp(tmp2, r, n, n2, ctx);    // r^n mod n^2
//     BN_mul(c, tmp, tmp2, ctx);          // c = g^m * r^n mod n^2
//     BN_mod(r2, r, n, ctx);              // r^2 mod n
//     BN_mod_exp(c2, c, n, n2, ctx);      // (1+c)^n mod n^2
//     BN_sub(c2, c2, BN_value_one());     // (1+c)^n - 1
//     BN_div(c2, NULL, c2, n, ctx);       // (1+c)^n - 1 / n
//     BN_mod_inverse(tmp, r2, n, ctx);         // r^{-1} mod n
//     BN_mul(c, c2, tmp, ctx);            // c * r^{-1} mod n
//     BN_free(n2);
//     BN_free(tmp);
//     BN_free(tmp2);
//     BN_free(r2);
//     BN_free(c2);
//     BN_CTX_free(ctx);
// }
//
// void paillier_decrypt(BIGNUM* m, BIGNUM* c, paillier_keypair* kp) {
//     BN_CTX* ctx = BN_CTX_new();
//     BIGNUM* n = kp->n;
//     BIGNUM* n2 = BN_new();
//     BN_mul(n2, n, n, ctx);  // n^2
//     BIGNUM* lambda = BN_new();
//     BIGNUM* u = BN_new();
//     BN_sub(lambda, n, BN_value_one());  // lambda = n-1
//     BN_div(lambda, NULL, lambda, n, ctx);   // lambda = lambda/n
//     BN_mod_exp(m, c, lambda, n2, ctx);  // c^{lambda} mod n^2
//     BN_sub(m, m, BN_value_one());       // m = (c^{lambda} - 1) / n
//     BN_div(m, NULL, m, n, ctx);
//     BN_mod_inverse(u, lambda, n, ctx);     // lambda^{-1} mod n
//     BN_mul(m, m, u, ctx);            // lambda^{-1} * m mod n
//     BN_free(n2);
//     BN_free(lambda);
//     BN_free(u);
//     BN_CTX_free(ctx);
// }
//
// int main() {
//     paillier_keypair kp;
//     BIGNUM* m = BN_new();
//     BIGNUM* c = BN_new();
//     BIGNUM* r = BN_new();
//     paillier_genkey(&kp, 1024);
//     BN_set_word(m, 42);
//     paillier_encrypt(c, m, r, &kp);
//     paillier_decrypt(m, c, &kp);
//     printf("m: %s\n", BN_bn2dec(m));
//     BN_free(c);
//     BN_free(m);
//     BN_free(r);
//     BN_free(kp.n);
//     BN_free(kp.g);
//     return 0;
// }